
from pymongo import UpdateOne
from pev3.factor.BaseFactor import BaseFactor
from pev3.util.StockUtil import getTradingDateStockCodes, getTradingDates
from pev3.data.DataModule import DataModule
from datetime import datetime

"""
实现规模因子的计算和保存
"""
class MarketScaleFactor(BaseFactor):

    def __init__(self):
        BaseFactor.__init__(self, name='market_scale')

    """
    计算指定时间段内所有股票的该因子的值，并保存到数据库中
    :param startDate:  开始时间
    :param endDate: 结束时间
    """
    def compute(self, startDate, endDate):
        dataModule = DataModule()

        # 如果没有指定日期范围，则默认为计算当前交易日的数据
        if startDate is None:
            startDate = datetime.now().strftime('%Y-%m-%d')

        if endDate is None:
            endDate = datetime.now().strftime('%Y-%m-%d')

        dates = getTradingDates(startDate, endDate)

        for date in dates:
            #  查询出股票在某一交易日的总股本
            dfBasics = dataModule.getStockBasic(date)

            if dfBasics.index.size == 0:
                continue

            # 将索引修改为code
            dfBasics.set_index(['code'], 1, inplace=True)

            # 查询出股票在某一个交易日的收盘价
            dfDailies = dataModule.getOneDayKData(autype=None, date=date)

            if dfDailies.index.size == 0:
                continue

            # 将索引修改为code
            dfDailies.set_index(['code'], 1, inplace=True)

            updateRequests = []
            for code in dfDailies.index:
                try:
                    # 股份
                    close = dfDailies.loc[code]['close']
                    # 总股本
                    totals = dfDailies.loc[code]['totals']
                    # 总市值
                    totalCapital = round(close * totals, 2)
                    print('%s, %s, mkt_cap: %15.2f' %
                          (code, date, totalCapital),
                          flush=True)
                    updateRequests.append(
                        UpdateOne(
                            {'code': code, 'date': date},
                            {'$set': {'code': code, 'date': date, self.name: totalCapital}}, upsert=True
                        )
                    )
                except:
                    print('计算规模因子时发生异常，股票代码：%s，日期：%s'
                          % (code, date),
                          flush=True)

            if len(updateRequests) > 0:
                result = self.collection.bulk_write(updateRequests, ordered=False)
                print('股票代码: %s, 因子: %s, 插入：%4d, 更新: %4d' %
                      (code, self.name, result.upserted_count, result.modified_count), flush=True)

if __name__ == '__main__':
    # 执行因子的提取任务
    MarketScaleFactor().compute('2018-10-01', '2018-10-22')